#pragma once
#include "BigInt.hpp"
#include "MD5.hpp"

BigInt StrToBigInt(const std::string &content)
{
    BigInt result;
    BigInt tmp = 1;
    for (size_t i = 0; i < content.length(); ++i)
    {
        result += tmp * ((int)content[i] + 128);
        tmp *= 256;
    }
    return result;
}

class ElGamal
{
public:
    ElGamal()
    {
        // 选择一个大素数 p
        // 这里的 p 不可以选择过大 因为计算原根的算法并不能保证大位是对的
        p_ = BigInt::GetRandomPrimeBigInt(48);
        // 选择一个原根 a
        a_ = p_.PrimitiveRoot();
        // 选择一个整数 d 作为私钥
        d_ = BigInt::GetRandomBigInt(20);
        // 计算 b = a 的 d 次方 mod p
        b_ = a_.powm(d_, p_);
        // 私钥 d     公钥 p a b
    }

    void Sign(std::string& content, BigInt& s1, BigInt& s2)
    {
        MD5 md;
        std::string source = md.encode(content);
        // 将 hash 转化为 数字
        BigInt hm = StrToBigInt(source);
        std::cout << "==============进行签名==============" << std::endl;
        std::cout << "签名时的hash值: " << hm << std::endl;
        BigInt k = BigInt::GetRandomBigInt(30);
        std::cout << "临时密钥k: " << k << std::endl;
        s1 = a_.powm(k, p_);
        s2 = BigInt::ModFraction(hm - d_ * s1, k, p_ - 1);
        std::cout << "签名对: (" << s1 << ", " << s2 << ")" << std::endl;
        std::cout << "==============签名结束==============" << std::endl << std::endl;
    }

    bool Verify(std::string& content, const BigInt& s1, const BigInt& s2)
    {
        MD5 md;
        std::string source = md.encode(content);
        // 将 hash 转化为 数字
        BigInt hm = StrToBigInt(source);
        std::cout << "==============进行验证==============" << std::endl;
        std::cout << "验证时的hash值: " << hm << std::endl;
        BigInt v1 = a_.powm(hm, p_);
        BigInt v2 = (b_.powm(s1, p_) * s1.powm(s2, p_)) % p_;
        std::cout << "验证生成: (" << v1 << ", " << v2 << ")" << std::endl;
        std::cout << (v1 == v2 ? "验证通过" : "验证不通过") << std::endl;
        std::cout << "==============验证结束==============" << std::endl << std::endl;
        return v1 == v2;
    }

    void Test()
    {
        std::cout << "本次ElGamal算法信息: 私钥: " << d_ << " 公钥: (" << a_ << ", " << b_ << ", " << p_ << ")" << std::endl;
        std::string content = "Nihao";
        std::cout << "输入待签名内容: ";
        std::cin >> content;
        BigInt s1;
        BigInt s2;
        Sign(content, s1, s2);
        Verify(content, s1, s2);
    }

private:
    BigInt p_;            // 素数
    BigInt a_;            // 原根
    BigInt b_;            // 公钥
    BigInt d_;            // 私钥
};
